from langchain_zhipu import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables import RunnablePassthrough
from operator import itemgetter

from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool

output_parser = StrOutputParser()

llm = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                  , model="glm-4")
llm.temperature = 0.01  # 温度设置为0，结果随机性 ghbnm
mysql_uri = 'mysql+pymysql://rdsroot:Geely%40db20211206@10.240.53.162:3306/test'

db = SQLDatabase.from_uri(mysql_uri)


def sqlresult(text):
    return text.split('\n')[1]


template = """Just answer SQL,exclude all information outside of SQL,Based on the table schema below, write a SQL query that would answer the user's question:
{schema}

Question: {question}
SQL Query:"""
prompt = ChatPromptTemplate.from_template(template)
db = SQLDatabase.from_uri(mysql_uri)


def get_schema(_):
    return db.get_table_info()


sql_response = (
        RunnablePassthrough.assign(schema=get_schema)
        | prompt
        | llm.bind(stop=["\nSQLResult:"])
        | output_parser
)

print(
    "question sql:" + sql_response.invoke(
        {"question": "which user's phone number is 13757136095,what's he delflag and emp_no?"}))
print("======================================")
execute_query = QuerySQLDataBaseTool(db=db)
chain = sql_response | execute_query
print("question answer:" + chain.invoke(
    {"question": "which user's phone number is 13757136095,what's he delflag  and emp_no?"}))

answer_prompt = ChatPromptTemplate.from_template(
    """Given the following user question, corresponding SQL query, and SQL result, answer the user question,
    In the database, tenant Id=22 represents 吉行生活, and null represents 吉好的. 用中文回答最终答案
    Question: {question}
    SQL Query: {query}
    SQL Result: {result}
    Answer: """
)

answer = answer_prompt | llm | StrOutputParser()

final_chain = (
        RunnablePassthrough.assign(query=sql_response).assign(result=itemgetter("query") | execute_query)
        | answer
)

print(final_chain.invoke({"question": "which user's phone number is 13757136095,what's he delflag  and emp_no and tenantId?"}))
